import { parseIntWithAutoBase, toBinaryString, toHexString } from './utils';
export interface Instruction {
  opcode: string;
  operands: string[];
}
export const opcodes = {
  MOV: '0100',
  ADD: '0000',
  SUB: '1000',
  AND: '0001',
  OR: '1001',
  RR: '1010',
  INC: '0111',
  DEC: '1011',
  LAD: '1100',
  STA: '1101',
  JMP: '1110',
  BZC: '1111',
  IN: '0010',
  OUT: '0011',
  LDI: '0110',
  HALT: '0101',
  STI: '0111',
  CLI: '1000',
  PUSH: '1001',
  POP: '1010',
  IRET: '1011',
} as const;
export function splitInstruction(instruction: string): Instruction {
  instruction = instruction.toUpperCase().trim();
  const [opcode, ...operandsPart] = instruction.split(' ').map((s) => s.trim());
  const operands = operandsPart
    .join(' ')
    .split(',')
    .map((s) => s.trim());
  return { opcode, operands };
}
export type Opcode = keyof typeof opcodes;
export function isOpcode(opcode: string): opcode is Opcode {
  return opcode in opcodes;
}
export function encodeOpcode(opcode: string) {
  if (isOpcode(opcode)) {
    return opcodes[opcode];
  }
  throw new Error(`Invalid opcode: ${opcode}`);
}
export const registerMap = {
  R0: '00',
  R1: '01',
  R2: '10',
  R3: '11',
  RI: '10',
  SP: '11',
} as const;
export type Register = keyof typeof registerMap;
export const isRegister = (register: string): register is Register =>
  register in registerMap;
export function encodeRegister(register: string) {
  if (isRegister(register)) {
    return registerMap[register];
  }
  throw new Error(`Invalid register: ${register}`);
}
export function encodeImmediate(immediate: string) {
  const immediateNumber = parseIntWithAutoBase(immediate);
  if (immediateNumber >= 0 && immediateNumber < 256) {
    return toBinaryString(immediateNumber);
  }
  throw new Error(`Invalid immediate: ${immediate}`);
}
export function encodeInstruction(instruction: Instruction): string[] {
  const { opcode, operands } = instruction;
  const opE = encodeOpcode(opcode);
  const breakOut = (): never => {
    throw new Error(`Invalid instruction: ${opcode} ${operands}`);
  };
  if (!isOpcode(opcode)) {
    return breakOut();
  }
  switch (opcode) {
    case 'MOV':
    case 'ADD':
    case 'SUB':
    case 'AND':
    case 'OR':
    case 'RR': {
      const [rd, rs] = operands;
      if (!rd || !rs) return breakOut();
      const rdE = encodeRegister(rd);
      const rsE = encodeRegister(rs);
      return [opE + rsE + rdE];
    }
    case 'POP':
    case 'INC':
    case 'DEC': {
      const [rd] = operands;
      if (!rd) return breakOut();
      const rdE = encodeRegister(rd);
      return [opE + '00' + rdE];
    }
    case 'PUSH': {
      const [rs] = operands;
      if (!rs) return breakOut();
      const rsE = encodeRegister(rs);
      return [opE + rsE + '00'];
    }
    case 'HALT':
    case 'STI':
    case 'CLI':
    case 'IRET': {
      return [opE + '0000'];
    }
    case 'LDI': {
      const [rd, d] = operands;
      if (!rd || !d) return breakOut();
      const rdE = encodeRegister(rd);
      const dE = encodeImmediate(d);
      return [opE + '00' + rdE, dE];
    }
    case 'JMP':
    case 'BZC': {
      const [label] = operands;
      if (!label) return breakOut();
      return [opE + '0000', 'LABEL ' + label];
    }
    case 'IN': {
      const [rd, p] = operands;
      if (!rd || !p) return breakOut();
      const rdE = encodeRegister(rd);
      const pE = encodeImmediate(p);
      return [opE + '00' + rdE, pE];
    }
    case 'OUT': {
      const [p, rs] = operands;
      if (!rs || !p) return breakOut();
      const rsE = encodeRegister(rs);
      const pE = encodeImmediate(p);
      return [opE + rsE + '00', pE];
    }
    case 'LAD':
    case 'STA': {
      const [rd, ...md] = operands;
      if (!rd || md.length == 0) return breakOut();
      const rdE = encodeRegister(rd);
      if (md.length == 1) {
        const [d] = md;
        const asPtr = d.startsWith('[') && d.endsWith(']');
        const dE = encodeImmediate(asPtr ? d.slice(1, -1) : d);
        const mE = asPtr ? '01' : '00';
        return [opE + mE + rdE, dE];
      } else {
        const [m, d] = md;
        const mE = m == '[RI]' ? '10' : m == '[PC]' ? '11' : breakOut();
        const dE = encodeImmediate(d);
        return [opE + mE + rdE, dE];
      }
    }
  }
}
export function encodeProgram(program: string): string {
  const lines = program
    .split('\n')
    .map((s) => s.toUpperCase().split(';')[0].trim())
    .filter((s) => s != '');
  const labels = new Map<string, number>();
  const encodedProgram = new Map<number, string>();
  const script = new Map<number, string[]>();
  let pos = 0;
  const addScript = (str: string) => {
    const scriptLine = script.get(pos);
    if (scriptLine != undefined) {
      script.set(pos, [...scriptLine, str]);
    } else {
      script.set(pos, [str]);
    }
  };
  for (const line of lines) {
    if (line.endsWith(':')) {
      const label = line.slice(0, -1);
      addScript(`(${label}:)`);
      labels.set(label, pos);
    } else if (line.startsWith('@')) {
      pos = parseIntWithAutoBase(line.slice(1));
    } else if (line.startsWith('=')) {
      addScript(line);
      encodedProgram.set(pos++, encodeImmediate(line.slice(1)));
    } else {
      const instruction = splitInstruction(line);
      const encodedInstruction = encodeInstruction(instruction);
      addScript(line);
      for (const encodedLine of encodedInstruction) {
        encodedProgram.set(pos++, encodedLine);
      }
    }
  }
  for (const [idx, line] of encodedProgram.entries()) {
    const prefix = 'LABEL ';
    if (line.startsWith(prefix)) {
      const label = line.slice(prefix.length);
      const labelPos = labels.get(label);
      if (labelPos != undefined) {
        encodedProgram.set(idx, toBinaryString(labelPos));
      } else {
        throw new Error(`Unknown label: ${label}`);
      }
    }
  }
  const result: string[] = [];
  pos = 0;
  for (const [idx, line] of encodedProgram.entries()) {
    if (idx != pos) {
      result.push(`@${toHexString(idx)}H`);
      pos = idx;
    }
    const posStr = toBinaryString(idx);
    const resultLine = `${posStr}: ${line}`;
    const scriptLine = script.get(idx);
    result.push(resultLine + (scriptLine ? '; ' + scriptLine.join(' ') : ''));
    pos++;
  }
  return result.join('\n');
}
